// Package network provides network optimization features for on-demand loading
package network

import (
	"io"
	"sync"
	"time"
)

// BandwidthLimiter provides bandwidth limiting for downloads
type BandwidthLimiter struct {
	limit         int64 // bytes per second
	bucket        *tokenBucket
	mu            sync.Mutex
}

// tokenBucket implements token bucket algorithm for rate limiting
type tokenBucket struct {
	tokens    int64
	capacity  int64
	refillRate int64
	lastRefill time.Time
	mu        sync.Mutex
}

// NewBandwidthLimiter creates a new bandwidth limiter
func NewBandwidthLimiter(limit int64) *BandwidthLimiter {
	if limit <= 0 {
		return &BandwidthLimiter{limit: 0} // No limit
	}

	return &BandwidthLimiter{
		limit:  limit,
		bucket: newTokenBucket(limit, limit),
	}
}

// newTokenBucket creates a new token bucket
func newTokenBucket(capacity, refillRate int64) *tokenBucket {
	return &tokenBucket{
		tokens:     capacity,
		capacity:   capacity,
		refillRate: refillRate,
		lastRefill: time.Now(),
	}
}

// Reader returns a bandwidth-limited reader
func (bl *BandwidthLimiter) Reader(r io.Reader) io.Reader {
	if bl.limit <= 0 {
		return r // No limit
	}
	return &limitedReader{
		reader:  r,
		limiter: bl,
	}
}

// Writer returns a bandwidth-limited writer
func (bl *BandwidthLimiter) Writer(w io.Writer) io.Writer {
	if bl.limit <= 0 {
		return w // No limit
	}
	return &limitedWriter{
		writer:  w,
		limiter: bl,
	}
}

// Wait waits for the specified number of bytes to be available
func (bl *BandwidthLimiter) Wait(bytes int64) error {
	if bl.limit <= 0 {
		return nil // No limit
	}

	return bl.bucket.consume(bytes)
}

// SetLimit updates the bandwidth limit
func (bl *BandwidthLimiter) SetLimit(limit int64) {
	bl.mu.Lock()
	defer bl.mu.Unlock()

	bl.limit = limit
	if limit <= 0 {
		bl.bucket = nil
	} else {
		bl.bucket = newTokenBucket(limit, limit)
	}
}

// GetLimit returns the current bandwidth limit
func (bl *BandwidthLimiter) GetLimit() int64 {
	bl.mu.Lock()
	defer bl.mu.Unlock()
	return bl.limit
}

// consume consumes tokens from the bucket
func (tb *tokenBucket) consume(tokens int64) error {
	tb.mu.Lock()
	defer tb.mu.Unlock()

	// Refill tokens based on elapsed time
	now := time.Now()
	elapsed := now.Sub(tb.lastRefill)
	tb.lastRefill = now

	// Calculate tokens to add
	tokensToAdd := int64(elapsed.Seconds() * float64(tb.refillRate))
	if tokensToAdd > 0 {
		tb.tokens = min64(tb.tokens+tokensToAdd, tb.capacity)
	}

	// Wait if not enough tokens
	if tb.tokens < tokens {
		needed := tokens - tb.tokens
		waitTime := time.Duration(float64(needed) / float64(tb.refillRate) * float64(time.Second))
		time.Sleep(waitTime)

		// Refill again after waiting
		now = time.Now()
		elapsed = now.Sub(tb.lastRefill)
		tb.lastRefill = now
		tokensToAdd = int64(elapsed.Seconds() * float64(tb.refillRate))
		if tokensToAdd > 0 {
			tb.tokens = min64(tb.tokens+tokensToAdd, tb.capacity)
		}
	}

	// Consume tokens
	if tb.tokens >= tokens {
		tb.tokens -= tokens
		return nil
	}

	return nil
}

// limitedReader implements bandwidth-limited reading
type limitedReader struct {
	reader  io.Reader
	limiter *BandwidthLimiter
}

// Read reads data with bandwidth limiting
func (lr *limitedReader) Read(p []byte) (int, error) {
	if lr.limiter.limit <= 0 {
		return lr.reader.Read(p)
	}

	// Read a limited amount of data
	maxRead := int64(len(p))
	if err := lr.limiter.Wait(maxRead); err != nil {
		return 0, err
	}

	n, err := lr.reader.Read(p)
	return n, err
}

// limitedWriter implements bandwidth-limited writing
type limitedWriter struct {
	writer  io.Writer
	limiter *BandwidthLimiter
}

// Write writes data with bandwidth limiting
func (lw *limitedWriter) Write(p []byte) (int, error) {
	if lw.limiter.limit <= 0 {
		return lw.writer.Write(p)
	}

	// Write in chunks to respect bandwidth limit
	totalWritten := 0
	chunkSize := int(lw.limiter.limit / 10) // Write in smaller chunks
	if chunkSize == 0 {
		chunkSize = 1024
	}

	for i := 0; i < len(p); i += chunkSize {
		end := i + chunkSize
		if end > len(p) {
			end = len(p)
		}

		chunk := p[i:end]
		if err := lw.limiter.Wait(int64(len(chunk))); err != nil {
			return totalWritten, err
		}

		n, err := lw.writer.Write(chunk)
		totalWritten += n
		if err != nil {
			return totalWritten, err
		}
	}

	return totalWritten, nil
}

// min64 returns the minimum of two int64 values
func min64(a, b int64) int64 {
	if a < b {
		return a
	}
	return b
}